import pandas as pd

from ReadData import getMatchHandledData, getSearchHandledData, getSearchOneDayData, pathName
import jieba
import pkuseg as ps
import thulac as tl

cutTypeCode = ["list(jieba.cut_for_search(sentence))", "list(jieba.cut(sentence, cut_all=True))",
               "ps_seg.cut(sentence)",
               "[item[0] for item in tl_seg.cut(sentence)]"]


def DivideSearchDirectory(cutType):
    for i in range(31):
        idx = i + 1
        if idx < 10:
            idx = "0" + str(idx)
        else:
            idx = str(idx)
        single_data = getSearchOneDayData(idx)
        if single_data is not None:
            divided_words = []
            search_words = single_data["word"]
            if cutType == 0:
                for words in search_words:
                    temp = []
                    words = words[1:-1]
                    divide_result = list(jieba.cut_for_search(words))
                    temp.extend(divide_result)
                    divided_words.append(temp)
            elif cutType == 1:
                for words in search_words:
                    temp = []
                    words = words[1:-1]
                    divide_result = list(jieba.cut(words, cut_all=True))
                    temp.extend(divide_result)
                    divided_words.append(temp)
            elif cutType == 2:
                seg = ps.pkuseg()
                for words in search_words:
                    temp = []
                    words = words[1:-1]
                    divide_result = seg.cut(words)
                    temp.extend(divide_result)
                    divided_words.append(temp)
            elif cutType == 3:
                seg = tl.thulac(seg_only=True)
                for words in search_words:
                    temp = []
                    words = words[1:-1]
                    divide_result = [item[0] for item in seg.cut(words)]
                    temp.extend(divide_result)
                    divided_words.append(temp)
            single_data["divide_result"] = divided_words
            single_data.to_csv("data/HandledData/" + pathName[cutType] + "/200608" + idx + ".csv", index=False)
            print("Success cutType:" + str(cutType) + " index" + idx)


def DivideMatchData(cutType):
    origin_data = getMatchHandledData()
    origin_data["SearchInfo"] = origin_data["SearchInfo"].astype(str)
    data = origin_data["SearchInfo"]
    result = []
    if cutType == 2:
        ps_seg = ps.pkuseg()
    elif cutType == 3:
        tl_seg = tl.thulac()
    for index, sentence in enumerate(data):
        temp = []
        divide_result = eval(cutTypeCode[cutType])
        temp.extend(divide_result)
        result.append(temp)
        print(pathName[cutType] + ":" + str(index))
    origin_data["divide_result"] = result
    return origin_data
